import os
from paddlespeech.cli.tts import TTSExecutor

# 假设你已经安装好了 paddlespeech: pip install paddlespeech

# 1. 剧本解析
def parse_script(script_path):
    roles_lines = {}
    with open(script_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line or ':' not in line:
                continue
            role, speech = line.split(':', 1)
            role = role.strip()
            speech = speech.strip()
            if role not in roles_lines:
                roles_lines[role] = []
            roles_lines[role].append(speech)
    return roles_lines

# 2. 角色到声音模型映射（你可以选择不同的 TTS voice）
ROLE_VOICE = {
    "小明": "zh_canton_female",    # 示例模型，实际用 PaddleSpeech 支持的模型
    "小红": "zh_canton_male",
    "旁白": "zh_zhinan_female"
}
DEFAULT_VOICE = "zh_zhinan_female"

# 3. 生成音频
def synthesize_lines(roles_lines, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    tts = TTSExecutor()
    audio_files = []
    idx = 1
    for role, lines in roles_lines.items():
        voice = ROLE_VOICE.get(role, DEFAULT_VOICE)
        for line in lines:
            # PaddleSpeech CLI TTS 支持指定 am/ voc/ speaker
            # 这里只用声学模型/声码器/发音人做示例，实际参数请参考 PaddleSpeech 文档
            output_file = os.path.join(output_dir, f"{idx:03d}_{role}.wav")
            tts(text=line, output=output_file, am=voice)
            audio_files.append(output_file)
            idx += 1
    return audio_files

# 4. 拼接音频（简单合成，有需要可用更专业音频库如 pydub）
from pydub import AudioSegment

def concatenate_audios(audio_files, output_file):
    final_audio = AudioSegment.empty()
    for file in audio_files:
        segment = AudioSegment.from_wav(file)
        final_audio += segment
    final_audio.export(output_file, format="wav")

if __name__ == "__main__":
    script_path = "drama.txt"   # 剧本文件
    output_dir = "audio_segments"
    final_output = "final_drama.wav"

    roles_lines = parse_script(script_path)
    audio_files = synthesize_lines(roles_lines, output_dir)
    concatenate_audios(audio_files, final_output)
    print(f"有声剧已生成: {final_output}")